"""
Data loading and preprocessing utilities.
"""
import pandas as pd
from typing import Optional, List, Dict, Any, Tuple
from datetime import datetime


class DataLoader:
    """Handles data loading and preprocessing for attribution analysis."""

    def __init__(self):
        self.data = None
        self.metadata = {}

    def load_from_dataframe(
        self,
        df: pd.DataFrame,
        metric_columns: List[str],
        dimension_columns: List[str],
        time_column: str = 'date'
    ) -> pd.DataFrame:
        """
        Load data from a pandas DataFrame.

        Args:
            df: Source DataFrame
            metric_columns: List of metric column names
            dimension_columns: List of dimension column names
            time_column: Name of the time column

        Returns:
            Processed DataFrame
        """
        self.data = df.copy()

        # Convert time column to datetime
        if time_column in self.data.columns:
            self.data[time_column] = pd.to_datetime(self.data[time_column])

        # Store metadata
        self.metadata = {
            'metric_columns': metric_columns,
            'dimension_columns': dimension_columns,
            'time_column': time_column,
            'total_rows': len(self.data)
        }

        return self.data

    def filter_by_time_range(
        self,
        df: pd.DataFrame,
        time_column: str,
        start_date: datetime,
        end_date: datetime
    ) -> pd.DataFrame:
        """Filter data by time range."""
        mask = (df[time_column] >= start_date) & (df[time_column] <= end_date)
        return df[mask].copy()

    def filter_by_dimensions(
        self,
        df: pd.DataFrame,
        filters: Dict[str, Any]
    ) -> pd.DataFrame:
        """
        Filter data by dimension values.

        Args:
            df: Source DataFrame
            filters: Dictionary of {dimension: value} filters

        Returns:
            Filtered DataFrame
        """
        result = df.copy()
        for dimension, value in filters.items():
            if isinstance(value, list):
                result = result[result[dimension].isin(value)]
            else:
                result = result[result[dimension] == value]
        return result

    def aggregate_by_dimensions(
        self,
        df: pd.DataFrame,
        dimensions: List[str],
        metrics: List[str],
        aggregation: str = 'sum'
    ) -> pd.DataFrame:
        """
        Aggregate data by specified dimensions.

        Args:
            df: Source DataFrame
            dimensions: List of dimension columns to group by
            metrics: List of metric columns to aggregate
            aggregation: Aggregation method ('sum', 'mean', 'count', etc.)

        Returns:
            Aggregated DataFrame
        """
        if not dimensions:
            # Aggregate all data
            if aggregation == 'sum':
                return pd.DataFrame({m: [df[m].sum()] for m in metrics})
            elif aggregation == 'mean':
                return pd.DataFrame({m: [df[m].mean()] for m in metrics})
            elif aggregation == 'count':
                return pd.DataFrame({m: [df[m].count()] for m in metrics})

        agg_dict = {m: aggregation for m in metrics}
        return df.groupby(dimensions, as_index=False).agg(agg_dict)

    def get_dimension_values(
        self,
        df: pd.DataFrame,
        dimension: str
    ) -> List[Any]:
        """Get unique values for a dimension."""
        return df[dimension].unique().tolist()

    def calculate_change_rate(
        self,
        current_value: float,
        baseline_value: float
    ) -> float:
        """Calculate percentage change between current and baseline."""
        if baseline_value == 0:
            return float('inf') if current_value > 0 else 0
        return (current_value - baseline_value) / baseline_value

    def split_by_period(
        self,
        df: pd.DataFrame,
        time_column: str,
        current_period: Tuple[datetime, datetime],
        baseline_period: Tuple[datetime, datetime]
    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Split data into current and baseline periods.

        Returns:
            Tuple of (current_df, baseline_df)
        """
        current_df = self.filter_by_time_range(
            df, time_column, current_period[0], current_period[1]
        )
        baseline_df = self.filter_by_time_range(
            df, time_column, baseline_period[0], baseline_period[1]
        )
        return current_df, baseline_df
